import os
import json

import numpy as np
import pandas as pd


if __name__ == "__main__":

    fn = "E:/workspace/SOMatch/tmp/os-select/res-select.xlsx"
    df = pd.read_excel(fn, sheet_name="Sheet1")

    path = np.array(df["path"]).tolist()
    
    def fn_sort(name):
        num = 0
        if "val" in name:
            num = 10000
        elif "test" in name:
            num = 20000

        num += int(name.split("opt")[1].split(".png")[0])

        return num

    path.sort(key=fn_sort)

    def get_sar(opt):
        sar = opt.replace("opt", "sar")
        return sar

    pair_path = list(map(get_sar, path))
    json_dict = {"OPT" : path, "SAR" : pair_path}
    train_dict = {"OPT" : path[:-500], "SAR" : pair_path[:-500]}
    val_dict = {"OPT" : path[-500:], "SAR" : pair_path[-500:]}

    with open("E:/workspace/SOMatch/tmp/json/os-select.json", "w") as fp:
        json.dump(json_dict, fp, indent="\t", sort_keys=True)
    with open("E:/workspace/SOMatch/tmp/json/os-s-train.json", "w") as fp:
        json.dump(train_dict, fp, indent="\t", sort_keys=True)
    with open("E:/workspace/SOMatch/tmp/json/os-s-val.json", "w") as fp:
        json.dump(val_dict, fp, indent="\t", sort_keys=True)
